"""
聊天模型调用工具, 虽然“工具调用”这个名称暗示模型直接执行某些操作，但实际上并非如此！模型仅生成工具的参数，实际运行工具（或不运行）取决于用户。
"""
from pprint import pprint
from typing import Dict, Callable, List

from langchain_core.tools import tool
from models import get_ds_model_client

@tool
def add(a: int, b: int):
    """add two numbers

    Args:
        a: first number,
        b: second number
    """
    return a +b

@tool
def multiply(a: int, b: int):
    """multiply two numbers

    Args:
        a: first number,
        b: second number
    """
    return a * b

def execute_tool(tool_calls_param: List[Dict]):
    """ 执行工具 """
    for tool_call in tool_calls_param:
        tool_name = tool_call['name']
        tool_args = tool_call['args']

        if tool_name in tool_map:
            # 调用对应的工具函数
            tool_function = tool_map[tool_name]
            result = tool_function.invoke(tool_args)
            print(f"{tool_name} result: {result}")
        else:
            print(f"Tool {tool_name} not found")

tools = [add, multiply]
tool_map: Dict[str, Callable] = {tool.name: tool for tool in tools}

llm = get_ds_model_client()
llm_bind_tools = llm.bind_tools(tools)

query = "What is 3 * 12?"

tool_calls = llm_bind_tools.invoke(query).tool_calls
# pprint(tool_calls)
# 调用工具函数
execute_tool(tool_calls)
